#include "tables.h"
#include "../config/config.h"
#include "../elf/elf.h"
#include "../debug/debug.h"
#include "../memory/malloc.h"
#include "../common/lib.h"
#include "../common/utils.h"
#include "../common/globals.h"
#include "../graphics/graphics.h"
#include "runtime_stubs.h"
#include <exploit_config.h>

#ifdef LOAD_MODULES_FOR_SYSCALLS
#ifndef AUTO_SEARCH_STUBS
#define AUTO_SEARCH_STUBS
#endif
#endif

// Adds NID entry to nid_table
int add_nid_to_table(u32 nid, u32 call, unsigned int lib_index)
{
	tGlobals * g = get_globals();
	NID_LOGSTR1("Adding NID 0x%08lX to table... ", nid);

	// Check if NID already exists in table (by another estimation for example)
	int index = get_nid_index(nid);

	// Doesn't exist, insert new
	if (index < 0)
	{
		if (g->nid_table.num >= NID_TABLE_SIZE)
		{
			exit_with_log("FATAL:NID TABLE IS FULL", NULL, 0);
			return -1;
		}

		g->nid_table.table[g->nid_table.num].nid = nid;
		g->nid_table.table[g->nid_table.num].call = call;
		g->nid_table.table[g->nid_table.num].lib_index = lib_index;		
		NID_LOGSTR1("Newly added @ %d\n", g->nid_table.num);
		g->nid_table.num++;
	}

	// If it exists, just change the old call with the new one
	else
	{
		g->nid_table.table[index].call = call;
		NID_LOGSTR1("Modified @ %d\n", g->nid_table.num);
	}

	return index;
}

// Adds a new library
int add_library_to_table(const tSceLibrary lib)
{
	tGlobals* g = get_globals();
	
	// Check if library exists
	int index = get_library_index(lib.name);

	// Doesn't exist, insert new
	if (index < 0)
	{
		// Check if there's space on the table
		if (g->library_table.num >= MAX_LIBRARIES)
		{
			LOGSTR1("->WARNING: Library table full, skipping new library %s\n", (u32)lib.name);
			return -1;
		}

		index = g->library_table.num;
        g->library_table.num++;
	}

	// Copying new library
	memcpy(&(g->library_table.table[index]), &lib, sizeof(lib));


	LOGSTR2("Added library %s @ %d\n", (u32)g->library_table.table[index].name, index);
	LOGLIB(g->library_table.table[index]);

	return index;
}

// Returns the index on nid_table for the given call
// NOTE: the whole calling instruction is to be passed
int get_call_index(u32 call)
{
	tGlobals * g = get_globals();
	int i = NID_TABLE_SIZE - 1;
	while(i >= 0 && g->nid_table.table[i].call != call)
	{
		i--;
	}

	return i;
}

// Gets i-th nid and its associated library
// Returns library name length
int get_lib_nid(u32 index, char* lib_name, u32* pnid)
{
	int ret;
	unsigned int num_libs, i = 0, count = 0;
	tImportedLibrary cur_lib;

	// DEBUG_PRINT("**GETTING NID INDEX:**", &index, sizeof(index));

	index += 1;
	ret = config_num_libraries(&num_libs);
	ret = config_first_library(&cur_lib);

	while (i<num_libs)
	{
		count += cur_lib.num_imports;
		if (index > count)
			ret = config_next_library(&cur_lib);
		else
		{
			ret = config_seek_nid(--index, pnid);
			// DEBUG_PRINT("**SEEK NID**", pnid, sizeof(u32));
			break;
		}
	}

	strcpy(lib_name, cur_lib.library_name);

	return strlen(lib_name);
}

// Return index in NID table for the call that corresponds to the NID pointed by "nid"
// Puts call in call_buffer
u32 get_call_nidtable(u32 nid, u32* call_buffer)
{
	tGlobals * g = get_globals();
	int i;
	
	*call_buffer = 0;
	for(i=0; i<NID_TABLE_SIZE; i++) 
	{
		if(nid == g->nid_table.table[i].nid)
		{
			if(call_buffer != NULL)
				*call_buffer = g->nid_table.table[i].call;
			break;
		}
	}

	return i;
}

// Return real instruction that makes the system call (jump or syscall)
u32 get_good_call(u32* call_pointer)
{
	// Dirty hack here :P but works
	if(*call_pointer & SYSCALL_MASK_IMPORT)
		call_pointer++;
	return *call_pointer;
}


int get_library_index_from_name(char* name)
{
	if (strcmp(name, "InterruptManager") == 0)
		return LIBRARY_INTERRUPTMANAGER;
	else if (strcmp(name, "ThreadManForUser") == 0)
		return LIBRARY_THREADMANFORUSER;
	else if (strcmp(name, "StdioForUser") == 0)
		return LIBRARY_STDIOFORUSER;
	else if (strcmp(name, "IoFileMgrForUser") == 0)
		return LIBRARY_IOFILEMGRFORUSER;
	else if (strcmp(name, "ModuleMgrForUser") == 0)
		return LIBRARY_MODULEMGRFORUSER;
	else if (strcmp(name, "SysMemUserForUser") == 0)
		return LIBRARY_SYSMEMUSERFORUSER;
	else if (strcmp(name, "sceSuspendForUser") == 0)
		return LIBRARY_SCESUSPENDFORUSER;
	else if (strcmp(name, "UtilsForUser") == 0)
		return LIBRARY_UTILSFORUSER;
	else if (strcmp(name, "LoadExecForUser") == 0)
		return LIBRARY_LOADEXECFORUSER;
	else if (strcmp(name, "sceGe_user") == 0)
		return LIBRARY_SCEGE_USER;
	else if (strcmp(name, "sceRtc") == 0)
		return LIBRARY_SCERTC;
	else if (strcmp(name, "sceAudio") == 0)
		return LIBRARY_SCEAUDIO;
	else if (strcmp(name, "sceDisplay") == 0)
		return LIBRARY_SCEDISPLAY;
	else if (strcmp(name, "sceCtrl") == 0)
		return LIBRARY_SCECTRL;
	else if (strcmp(name, "sceHprm") == 0)
		return LIBRARY_SCEHPRM;
	else if (strcmp(name, "scePower") == 0)
		return LIBRARY_SCEPOWER;
	else if (strcmp(name, "sceOpenPSID") == 0)
		return LIBRARY_SCEOPENPSID;
	else if (strcmp(name, "sceWlanDrv") == 0)
		return LIBRARY_SCEWLANDRV;
	else if (strcmp(name, "sceUmdUser") == 0)
		return LIBRARY_SCEUMDUSER;
	else if (strcmp(name, "sceUtility") == 0)
		return LIBRARY_SCEUTILITY;
	else if (strcmp(name, "sceSasCore") == 0)
		return LIBRARY_SCESASCORE;
	else if (strcmp(name, "sceImpose") == 0)
		return LIBRARY_SCEIMPOSE;
	else if (strcmp(name, "sceReg") == 0)
		return LIBRARY_SCEREG;
	else if (strcmp(name, "sceChnnlsv") == 0)
		return LIBRARY_SCECHNNLSV;
	else
		return -1;
}

// Returns index of NID in table
int get_nid_index(u32 nid)
{
	tGlobals * g = get_globals();
	int i;

	for (i=0; i<NID_TABLE_SIZE; i++)
	{
		if (g->nid_table.table[i].nid == 0)
			break;
		else if (g->nid_table.table[i].nid == nid)
			return i;
	}

	return -1;
}

// Checks if a library is in the global library description table
// Returns index if it's there
int get_library_index(const char* library_name)
{
	tGlobals * g = get_globals();
	LOGSTR1("Searching for library %s\n", (u32)library_name);
	if (library_name == NULL)
		return -1;

	int i;
	for (i=0; i<(int) g->library_table.num; i++)
	{
		//LOGSTR1("Current library: %s\n", g->library_table.table[i].library_name);
		if (g->library_table.table[i].name == NULL)
			break;
		else if (strcmp(library_name, g->library_table.table[i].name) == 0)
			return i;
	}

	return -1;
}

/*
 * Retrieves highest known syscall of the previous library,
 * and lowest known syscall of the next library, to get some
 * rough boundaries of where current library's syscalls should be
 * returns 1 on success, 0 on failure
 */

/* Fills NID Table */
/* Returns NIDs resolved */
/* "pentry" points to first stub header in game */
int build_nid_table()
{
	tGlobals * g = get_globals();
	int library_index = -1;
	u32 i = 0, j, k = 0;
	unsigned int nlib_stubs;
	u32 *cur_nid, *cur_call, syscall_num, good_call, nid;
	tStubEntry *pentry;	

#ifdef AUTO_SEARCH_STUBS
#ifdef LOAD_MODULES_FOR_SYSCALLS
	load_modules_for_stubs();
#endif

	u32 stubs[MAX_RUNTIME_STUB_HEADERS];
	unsigned int current_stubid = 0;
	nlib_stubs = (unsigned int) (search_stubs(stubs));
	LOGSTR1("Found %d stubs\n", nlib_stubs);
	if (nlib_stubs == 0)
		exit_with_log(" ERROR: NO LIBSTUBS DEFINED IN CONFIG ", NULL, 0);    
	CLEAR_CACHE;
	pentry = (tStubEntry*) stubs[current_stubid];
#else        
	int aux1, aux2;
	u32 aux3;   

	// Getting game's .lib.stub address
	if ((aux1 = config_initialize()) < 0)
		exit_with_log(" ERROR INITIALIZING CONFIG ", &aux1, sizeof(aux1));

	if ((aux1 = config_num_lib_stub(&nlib_stubs)) < 0)
		exit_with_log(" ERROR READING NUMBER OF LIBSTUBS FROM CONFIG ", &aux1, sizeof(aux1));

	if (nlib_stubs == 0)
		exit_with_log(" ERROR: NO LIBSTUBS DEFINED IN CONFIG ", NULL, 0);

	if ((aux2 = config_first_lib_stub(&aux3)) < 0)
		exit_with_log(" ERROR GETTING FIRST LIBSTUB FROM CONFIG ", &aux2, sizeof(aux2));

	pentry = (tStubEntry*) aux3;

#endif   
    
    
	//DEBUG_PRINT(" build_nid_table() ENTERING MAIN LOOP ", NULL, 0);

	do
	{
		NID_LOGSTR1("-->CURRENT MODULE LIBSTUB: 0x%08lX\n", (u32)pentry);
		
		// While it's a valid stub header
		while (elf_check_stub_entry(pentry))
		{
			if ((pentry->import_flags != 0x11) && (pentry->import_flags != 0))
			{
				// Variable import, skip it
				pentry = (tStubEntry*)((int)pentry + sizeof(u32));
				continue;
			}

			NID_LOGSTR1("\n-->Processing library: %s\n", (u32)(pentry->library_name));

			// Get actual call
			cur_call = pentry->jump_pointer;
			good_call = get_good_call(cur_call);

			// Only process if syscall and if the import is already resolved
			if (!(good_call & SYSCALL_MASK_RESOLVE) && (GET_SYSCALL_NUMBER(good_call) != SYSCALL_IMPORT_NOT_RESOLVED_YET))
			{		
				// Get current NID
				cur_nid = pentry->nid_pointer;

				// Is this library on the table?
				library_index = get_library_index((char *)pentry->library_name);

				// New library
				if (library_index < 0)
				{
					NID_LOGSTR1(" --> New: %d\n", k);
			
					strcpy(g->library_table.table[k].name, (char *)pentry->library_name);
					g->library_table.table[k].calling_mode = SYSCALL_MODE;
			
					// Get number of syscalls imported
					g->library_table.table[k].num_known_exports = pentry->stub_size;

					NID_LOGSTR1("Total known exports: %d\n", g->library_table.table[k].num_known_exports);
			
					// Browse all stubs defined by this header
					for(j=0; j<pentry->stub_size; j++)
					{
						nid = *cur_nid;
						NID_LOGSTR1("--Current NID: 0x%08lX\n", nid);
				
						// If NID is already in, don't put it again 
						if (get_nid_index(nid) < 0)
						{
							// Fill NID table																			
							add_nid_to_table(nid, get_good_call(cur_call), k);
							NID_LOGSTR1(" --> new inserted @ %d", i);

							// Check lowest syscall
							syscall_num = GET_SYSCALL_NUMBER(g->nid_table.table[i].call);
							NID_LOGSTR1(" with syscall 0x%08lX\n", syscall_num);
		                    
							i++;
						}
						cur_nid++;
						cur_call += 2;
					}

					// New library
					g->library_table.num++;

					// Next library entry
					k++;
				}

				// Old library
				else
				{
					NID_LOGSTR1(" --> Old: %d\n", library_index);

					LOGLIB(g->library_table.table[library_index]);
					
					NID_LOGSTR1("Number of imports of this stub: %d\n", pentry->stub_size);

					if (g->library_table.table[library_index].calling_mode != SYSCALL_MODE)
					{
						config_close();
						exit_with_log(" ERROR OLD CALL MODE IS SYSCALL, NEW IS JUMP ", &library_index, sizeof(library_index));
					}
			
					// Browse all stubs defined by this header
					for(j=0; j<pentry->stub_size; j++)
					{
						nid = *cur_nid;
						NID_LOGSTR1("--Current NID: 0x%08lX\n", nid);
				
						// If NID is already in, don't put it again 
						if (get_nid_index(nid) < 0)
						{
							// Fill NID table
							
// There is probably a way to shorten this #ifdef block, but I'm too lazy to check if I broke backwards compatibility here,
// so I'm being concervative(wololo)																			
							add_nid_to_table(nid, get_good_call(cur_call), library_index);
							NID_LOGSTR1(" --> new inserted @ %d", i);
				
							// Check lowest syscall
							syscall_num = GET_SYSCALL_NUMBER(g->nid_table.table[i].call);
							NID_LOGSTR1(" with syscall 0x%08lX\n", syscall_num);
							i++;
							g->library_table.table[library_index].num_known_exports++;
						}
						cur_nid++;
						cur_call += 2;
					}
				}
			}

			// Next entry
			pentry++;

			CLEAR_CACHE;
		}
	
#ifdef AUTO_SEARCH_STUBS
        current_stubid++;
 		// Next .lib.stub
		if (current_stubid < nlib_stubs)
		{
			pentry = (tStubEntry*) stubs[current_stubid];
		}
		else
			break;               
#else	
		nlib_stubs--;

		// Next .lib.stub
		if (nlib_stubs > 0)
		{
			aux2 = config_next_lib_stub(&aux3);
			if (aux2 < 0)
				exit_with_log(" ERROR GETTING NEXT LIBSTUB FROM CONFIG ", &aux2, sizeof(aux2));
			pentry = (tStubEntry*) aux3;
		}
		else
			break;
#endif
		sceKernelDelayThread(100000);
		
	} while(1);

	CLEAR_CACHE;

#ifdef AUTO_SEARCH_STUBS
#ifdef LOAD_MODULES_FOR_SYSCALLS
	unload_modules_for_stubs();
#endif
#else
	config_close();  
#endif
	
	return i;
}
